from fastapi import status
from fastapi.responses import JSONResponse
from typing import Union, Optional

from pydantic import BaseModel


class ErrorBase(BaseModel):
    code: int
    message: str = ""


def success_dict(data: Union[list, dict, str, int] = None, message: str = "Success"):
    return {
        "code": 200,
        "message": message,
        "data": data
    }


def success(data: Union[list, dict, str, int] = None, message: str = "Success"):
    """ 接口成功返回 """
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content=success_dict(data, message)
    )


def error_dict(error_base: ErrorBase, message: Optional[str] = None, message_append: str = "",
               data: Union[list, dict, str] = None):
    """ 错误接口返回 """
    return {
        "code": error_base.code,
        "message": (message or error_base.message) + message_append,
        "data": data
    }


def error(
        error_base: ErrorBase,
        message: Optional[str] = None,
        message_append: str = "",
        data: Union[list, dict, str] = None,
        status_code: int = status.HTTP_500_INTERNAL_SERVER_ERROR
):
    """ 错误接口返回 """
    return JSONResponse(
        status_code=status_code,
        content=error_dict(error_base, message, message_append, data)
    )
